#!/usr/bin/env python
import roslib; roslib.load_manifest('stack_objects')
import rospy
from geometry_msgs.msg import PointStamped
from pr2_tasks.tasks import Tasks
from pr2_tasks.pickplace import PickPlace
from pr2_tasks.tableware_detection import TablewareDetection
from pr2_tasks.pickplace_definitions import PlaceGoal
from pr2_python import base
from pr2_python.world_interface import WorldInterface
from pr2_python.planning_scene_interface import get_planning_scene_interface
from pr2_python.arm_mover import ArmMover
from pr2_tasks.pickup_tableware import PickupTableware
import copy

class ObjectStacker():
  def __init__(self):
    wi = WorldInterface()
    global res
    global ordered_objects

  def detect_objects(self):
     global res
     det = TablewareDetection()
     res = det.detect_objects()
     rospy.loginfo('Detected objects with the following labels and poses:')
     return

  def object_diameter(self, pg):
     y_max = pg.target.cluster.points[0].y
     y_min = pg.target.cluster.points[0].y
     for point in pg.target.cluster.points:
        if point.y > y_max:
           y_max = point.y
        if point.y < y_min:
           y_min = point.y
     return (y_max-y_min)

  def separate_graspable_objects(self):
     global res      
     for pg in res.pickup_goals:
        print pg.label
        if pg.label is not 'graspable':
           pass
        else:
           diam = object_diameter(pg)
           print diam
           if (diam < 0.07):
              pg.label = 'small_cups'
           elif (diam < 0.14):
              pg.label = 'cups'
           elif (diam < 0.2):
              pg.label = 'bowls'
           elif (diam < 0.35):
              pg.label = 'plates'
           else:
              pg.label = 'tablet'

  def create_stack_list(self):
     global res
     global ordered_objects
     ordered_objects = []
     known_objects = ['plate', 'bowl', 'cup']
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18980: #big sushi plate
              ordered_objects.append(pg)
        if pg.label == 'plates':
           ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18979: #little sushi plate
              ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18976: #big sushi bowl
              ordered_objects.append(pg)
        if pg.label == 'bowls':
           ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18975: #little sushi bowl
              ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18978: #big sushi cup
              ordered_objects.append(pg)
        if pg.label == 'cups':
           ordered_objects.append(pg)
     for pg in res.pickup_goals:
        if pg.label in known_objects:
           if pg.target.potential_models[0].model_id==18977: #little sushi cup
              ordered_objects.append(pg)
        if pg.label == 'small_cups':
           ordered_objects.append(pg)
     for pg in ordered_objects:
        rospy.loginfo('Label: '+str(pg.label))

  def exec_stack_objects(self):
     global ordered_objects
     searchplace = []
     for z in [0.02, 0.04, 0.08, 0.1]:
        #for a in [0.0, np.pi/2.0, -np.pi/2.0, np.pi]:
        for a in [0.0, 1.57, -1.57, -3.14]:
           searchplace.append((0, 0, z, a))

     pickplace = PickPlace(search_place = searchplace)

     if len(ordered_objects) == 0:
        print "Sorry no know objects found, returning"
        return

     if len(ordered_objects) == 1:
        print "Only one object found, only need to pick that one"
        return

     global put_down_place
     put_down_place = copy.deepcopy(ordered_objects[0].object_pose_stamped)

     for i in range(1,len(ordered_objects)):
        ordered_objects[i].arm_name = 'left_arm'
        # try to pick the object up
        try:
           pickplace.pickup(ordered_objects[i])
        except:
           print "Didn't succeed in picking up the object"

        print ordered_objects[i-1].collision_support_surface_name 
        # Generate the put down position
        place_pose_stamped = copy.deepcopy(put_down_place)
        place_goal = PlaceGoal(ordered_objects[i].arm_name, 
                               [place_pose_stamped], desired_approach_distance=0.15,
                               collision_support_surface_name = ordered_objects[i-1].collision_object_name,
                               collision_object_name = ordered_objects[i].collision_object_name)

        succes = 0
        #while succes == 0:
           #and put it back down
        try:
           pickplace.place(place_goal)
           succes = 1
        except:
           print "Do something smarter here please"
           print "Shoot missed it, can't put it down here"

  def exec_pose_stack_objects(self, pose):
     global ordered_objects
     searchplace = []
     for z in [0.02, 0.04, 0.08, 0.1]:
        #for a in [0.0, np.pi/2.0, -np.pi/2.0, np.pi]:
        for a in [0.0, 1.57, -1.57, -3.14]:
           searchplace.append((0, 0, z, a))

     pickplace = PickPlace(search_place = searchplace)

     if len(ordered_objects) == 0:
        print "Sorry no know objects found, returning"
        return

     #global put_down_place
     #put_down_place = copy.deepcopy(ordered_objects[0].object_pose_stamped)
     #put_down_place.object_pose_stamped.pose.position.x = pose.x
     #put_down_place.object_pose_stamped.pose.position.y = pose.y
     #put_down_place.object_pose_stamped.pose.position.z = pose.z

     for i in range(0,len(ordered_objects)):
        ordered_objects[i].arm_name = 'left_arm'
        # try to pick the object up
        try:
           pickplace.pickup(ordered_objects[i])
        except:
           print "Didn't succeed in picking up the object"

        print ordered_objects[i-1].collision_support_surface_name 
        # Generate the put down position
        #place_pose_stamped = copy.deepcopy(put_down_place)
        place_goal = PlaceGoal(ordered_objects[i].arm_name, 
                               [pose], desired_approach_distance=0.15,
                               collision_support_surface_name = ordered_objects[i-1].collision_object_name,
                               collision_object_name = ordered_objects[i].collision_object_name)

        succes = 0
        #while succes == 0:
           #and put it back down
        try:
           pickplace.place(place_goal)
           succes = 1
        except:
           print "Do something smarter here please"
           print "Shoot missed it, can't put it down here"

  def move_arms_back(self):
     tasks = Tasks()
     tasks.move_arms_to_side()

  def pose_stack_objects(self, pose):
     # reset the world
     wi = WorldInterface()
     wi.reset()
     psi = get_planning_scene_interface()
     psi.reset()

     self.move_arms_back()
     self.detect_objects()
     # see if we can make sense of the graspable objects
     #separate_graspable_objects()
     self.create_stack_list()

     psi.reset()
     rospy.sleep(0.5)
     print "Press a key when you're ready to start stacking"
     raw_input()

     # Finally stack them
     self.exec_pose_stack_objects(pose)
     self.move_arms_back()

  def stack_objects(self):
     # reset the world
     wi = WorldInterface()
     wi.reset()
     psi = get_planning_scene_interface()
     psi.reset()

     self.move_arms_back()
     self.detect_objects()
     # see if we can make sense of the graspable objects
     #separate_graspable_objects()
     self.create_stack_list()

     psi.reset()
     rospy.sleep(0.5)
     print "Press a key when you're ready to start stacking"
     raw_input()

     # Finally stack them
     self.exec_stack_objects()
     self.move_arms_back()
